#include "CTcpSocketImpl.h"
#include "CNetService.h"
#include "CNetServiceImpl.h"
#include "CUtil.h"

namespace TB
{
    CTcpSocketImpl::CTcpSocketImpl(int nSocketFD):
        m_nSocketFD(nSocketFD)
    {
        if (m_nSocketFD < 0)
            m_nSocketFD = INVALID_SOCKET;
        else
            SetNonBlock(m_nSocketFD);
    }

	CTcpSocketImpl::CTcpSocketImpl()
    {
	}

	CTcpSocketImpl::~CTcpSocketImpl()
	{
        Close();
	}

	bool CTcpSocketImpl::Init(const std::weak_ptr<CNetService>& poNetService, const std::weak_ptr<CTcpSocketImpl>& poSelf)
	{
        if (!m_poNetService.expired())
            return false;
        m_poNetService = poNetService;
        m_poSelf = poSelf;
        return true;
	}

    ENetErrorCode CTcpSocketImpl::Bind(UINT32 dwIp, UINT16 wPort)
    {
        if (m_nSocketFD == INVALID_SOCKET)
        {
            auto nSocketFD = CreateSocket();
            if (nSocketFD == INVALID_SOCKET)
            {
                return GetLastError();
            }
            m_nSocketFD = nSocketFD;
            SetNonBlock(m_nSocketFD);
        }
        struct sockaddr_in address = {};
        address.sin_family = AF_INET;
        address.sin_port = htons(wPort);
        address.sin_addr.s_addr = htonl(dwIp);
        if (0 != bind(m_nSocketFD, (sockaddr*)&address, sizeof(address)))
            return GetLastError();
        return ENetErrorCode::Success;
    }

    void CTcpSocketImpl::Connect(UINT32 dwIp, UINT16 wPort, UINT64 qwTimeoutMilliSeconds, STaskProcessor<ENetErrorCode>& sTaskProcessor)//主机字节序
	{
        if (IsConnecting())
        {
            sTaskProcessor.Resolve(ENetErrorCode::AlreadyInProgress);
            return;
        }
        if (IsConnected())
        {
            sTaskProcessor.Resolve(ENetErrorCode::IsConnected);
            return;
        }
        if (m_nSocketFD == INVALID_SOCKET)
        {
            auto nSocketFD = CreateSocket();
            if (nSocketFD == INVALID_SOCKET)
            {
                sTaskProcessor.Resolve(GetLastError());
                return;
            }
            m_nSocketFD = nSocketFD;
            SetNonBlock(m_nSocketFD);
        }
        struct sockaddr_in address = {};
        address.sin_family = AF_INET;
        address.sin_port = htons(wPort);
        address.sin_addr.s_addr = htonl(dwIp);
#if _WIN32
        int nRet = WSAConnect(m_nSocketFD, (sockaddr*)&address, sizeof(address), NULL, NULL, NULL, NULL);
#else
        int nRet = connect(m_nSocketFD, (sockaddr*)&address, sizeof(address));
#endif
        if (nRet == 0)
        {
            sTaskProcessor.Resolve(ENetErrorCode::Success);
            return;
        }
        else
        {
            ENetErrorCode eErrorCode = GetLastError();
            if (eErrorCode != ENetErrorCode::InProgress && eErrorCode != ENetErrorCode::WouldBlock)
            {
                sTaskProcessor.Resolve(eErrorCode);
                return;
            }
        }
        m_sConnectTask = sTaskProcessor;
        if (false == _AddPoll(true))
        {
            m_sConnectTask = STaskProcessor<ENetErrorCode>();
            sTaskProcessor.Resolve(ENetErrorCode::OperationAborted);
            return;
        }
        auto& poTimerContext = CCurThreadTimerContext::GetCurrent().GetTimerContext();
        if (poTimerContext == nullptr)
        {
            CThreadPool::StartTimer(qwTimeoutMilliSeconds, [sTaskProcessor]() mutable
            {
                sTaskProcessor.Resolve(ENetErrorCode::TimedOut);
            });
        }
        else
        {
            m_cbConnectTimeoutTimer = poTimerContext->StartTimer(qwTimeoutMilliSeconds, [sTaskProcessor]() mutable
            {
                sTaskProcessor.Resolve(ENetErrorCode::TimedOut);
            });
        }
	}

    void CTcpSocketImpl::Send(const std::pair<UINT8*, UINT32>& sBuffer, STaskProcessor<ENetErrorCode>& sTaskProcessor)
    {
        auto& roBack = m_oSendTaskQueue.emplace_back(1, sTaskProcessor);
        roBack.m_sBuffer[0] = sBuffer;
        if (m_oSendTaskQueue.size() == 1)
            _Send();
    }

    void CTcpSocketImpl::Send(const std::pair<UINT8*, UINT32>& sBuffer1, const std::pair<UINT8*, UINT32>& sBuffer2, STaskProcessor<ENetErrorCode>& sTaskProcessor)
    {
        auto& roBack = m_oSendTaskQueue.emplace_back(2, sTaskProcessor);
        roBack.m_sBuffer[0] = sBuffer1;
        roBack.m_sBuffer[1] = sBuffer2;
        if (m_oSendTaskQueue.size() == 1)
            _Send();
    }

    void CTcpSocketImpl::Send(const std::pair<UINT8*, UINT32>& sBuffer1, const std::pair<UINT8*, UINT32>& sBuffer2, const std::pair<UINT8*, UINT32>& sBuffer3, STaskProcessor<ENetErrorCode>& sTaskProcessor)
    {
        auto& roBack = m_oSendTaskQueue.emplace_back(3, sTaskProcessor);
        roBack.m_sBuffer[0] = sBuffer1;
        roBack.m_sBuffer[1] = sBuffer2;
        roBack.m_sBuffer[2] = sBuffer3;
        if (m_oSendTaskQueue.size() == 1)
            _Send();
    }

    void CTcpSocketImpl::Send(const std::pair<UINT8*, UINT32>* psBuffers, UINT32 dwCount, STaskProcessor<ENetErrorCode>& sTaskProcessor)
    {
        auto& roBack = m_oSendTaskQueue.emplace_back(dwCount, sTaskProcessor);
        for (UINT32 i = 0; i != dwCount; ++i)
            roBack.m_sBuffer[i] = psBuffers[i];
        if (m_oSendTaskQueue.size() == 1)
            _Send();
    }

    void CTcpSocketImpl::_Send()
    {
        while (!m_oSendTaskQueue.empty())
        {
            SSendInfo& sSendInfo = m_oSendTaskQueue.front();
            for (; sSendInfo.m_dwToSendIndex < sSendInfo.m_sBuffer.size(); ++sSendInfo.m_dwToSendIndex)
            {
                std::pair<UINT8*, UINT32>& rsBuffer = sSendInfo.m_sBuffer[sSendInfo.m_dwToSendIndex];
                int nRet = -1;
                ENetErrorCode eNetErrorCode = ENetErrorCode::Success;
                do
                {
                    nRet = send(m_nSocketFD, (char*)rsBuffer.first, (int)rsBuffer.second, 0);
                    if (nRet > 0)
                    {
                        rsBuffer.first += (UINT32)nRet;
                        rsBuffer.second -= (UINT32)nRet;
                    }
                } while ((nRet > 0 && rsBuffer.second != 0) || (nRet < 0 && (eNetErrorCode = GetLastError()) == ENetErrorCode::Interrupted));
                if (nRet < 0 && eNetErrorCode != ENetErrorCode::TryAgain && eNetErrorCode != ENetErrorCode::WouldBlock)
                {
                    _RemovePollOut();
                    while (!m_oSendTaskQueue.empty())
                    {
                        SSendInfo& sSendInfo1 = m_oSendTaskQueue.front();
                        sSendInfo1.m_oProcessor.Resolve(eNetErrorCode);
                        m_oSendTaskQueue.pop_front();
                    }
                    return;
                }
                if (rsBuffer.second != 0)
                    break;
            }
            if (sSendInfo.m_dwToSendIndex < sSendInfo.m_sBuffer.size())
            {
                if (!_AddPoll(true))
                {
                    while (!m_oSendTaskQueue.empty())
                    {
                        SSendInfo& sSendInfo1 = m_oSendTaskQueue.front();
                        sSendInfo1.m_oProcessor.Resolve(ENetErrorCode::OperationAborted);
                        m_oSendTaskQueue.pop_front();
                    }
                }
                return;
            }
            sSendInfo.m_oProcessor.Resolve(ENetErrorCode::Success);
            m_oSendTaskQueue.pop_front();
        }
#if _WIN32
        _RemovePollOut();
#endif
    }

    void CTcpSocketImpl::Recv(const std::pair<UINT8*, UINT32>& sBuffer, STaskProcessor<std::pair<ENetErrorCode, UINT32>>& sTaskProcessor)
    {
        m_oRecvTaskQueue.emplace_back(false, sBuffer, sTaskProcessor);
        if (m_oRecvTaskQueue.size() == 1)
            _Recv();
    }

    void CTcpSocketImpl::RecvSome(const std::pair<UINT8*, UINT32>& sBuffer, STaskProcessor<std::pair<ENetErrorCode, UINT32>>& sTaskProcessor)
    {
        m_oRecvTaskQueue.emplace_back(true, sBuffer, sTaskProcessor);
        if (m_oRecvTaskQueue.size() == 1)
            _Recv();
    }

    void CTcpSocketImpl::_Recv()
    {
        while (!m_oRecvTaskQueue.empty())
        {
            SRecvInfo& sRecvInfo = m_oRecvTaskQueue.front();
            std::pair<UINT8*, UINT32>& rsBuffer = sRecvInfo.m_sBuffer;
            int nRet = -1;
            ENetErrorCode eNetErrorCode = ENetErrorCode::Success;
            do
            {
                nRet = recv(m_nSocketFD, (char*)rsBuffer.first, (int)rsBuffer.second, 0);
                if (nRet > 0)
                {
                    rsBuffer.first += (UINT32)nRet;
                    rsBuffer.second -= (UINT32)nRet;
                    sRecvInfo.m_dwRecvCount += (UINT32)nRet;
                }
            } while ((nRet > 0 && rsBuffer.second != 0 && !sRecvInfo.m_bIsRecvSome) || (nRet < 0 && (eNetErrorCode = GetLastError()) == ENetErrorCode::Interrupted));
            if (eNetErrorCode == ENetErrorCode::TryAgain || eNetErrorCode == ENetErrorCode::WouldBlock)
            {
                if (!_AddPoll(false))
                {
                    while (!m_oRecvTaskQueue.empty())
                    {
                        SRecvInfo& sRecvInfo1 = m_oRecvTaskQueue.front();
                        sRecvInfo1.m_oProcessor.Resolve({ ENetErrorCode::OperationAborted, 0 });
                        m_oRecvTaskQueue.pop_front();
                    }
                }
                return;
            }
            if (eNetErrorCode != ENetErrorCode::Success)
            {
                _RemovePollIn();
                while (!m_oRecvTaskQueue.empty())
                {
                    SRecvInfo& sRecvInfo1 = m_oRecvTaskQueue.front();
                    sRecvInfo1.m_oProcessor.Resolve({ eNetErrorCode, 0 });
                    m_oRecvTaskQueue.pop_front();
                }
                return;
            }
            sRecvInfo.m_oProcessor.Resolve({ ENetErrorCode::Success, sRecvInfo.m_dwRecvCount });
            m_oRecvTaskQueue.pop_front();
        }
#if _WIN32
        _RemovePollIn();
#endif
    }

    bool CTcpSocketImpl::GetLocalIpPort(std::pair<UINT32, UINT16>& roIpPort) const //主机字节序
    {
        if (m_nSocketFD == INVALID_SOCKET)
        {
            return false;
        }
        struct sockaddr_in addr;
#if _WIN32
        int addr_len = sizeof(addr);
#else
        socklen_t addr_len = sizeof(addr);
#endif
        if (0 == getsockname(m_nSocketFD, (struct sockaddr*)&addr, &addr_len))
        {
            roIpPort = { (UINT32)ntohl(addr.sin_addr.s_addr), (UINT16)ntohs(addr.sin_port) };
            return true;
        }
        return false;
    }

    bool CTcpSocketImpl::GetRemoteIpPort(std::pair<UINT32, UINT16>& roIpPort) const //主机字节序
    {
        if (m_nSocketFD == INVALID_SOCKET)
        {
            return false;
        }
        struct sockaddr_in addr;
#if _WIN32
        int addr_len = sizeof(addr);
#else
        socklen_t addr_len = sizeof(addr);
#endif
        if (0 == getpeername(m_nSocketFD, (struct sockaddr*)&addr, &addr_len))
        {
            roIpPort = { (UINT32)ntohl(addr.sin_addr.s_addr), (UINT16)ntohs(addr.sin_port) };
            return true;
        }
        return false;
    }

	bool CTcpSocketImpl::IsClosed() const
	{
        return m_nSocketFD == INVALID_SOCKET;
	}

	void CTcpSocketImpl::Close()
	{
        if (IsClosed())
            return;
        if (m_nSocketFD != INVALID_SOCKET)
        {
            if (m_nSettedEvents != 0)
            {
                auto poNetService = m_poNetService.lock();
                if (poNetService && 0 == poNetService->GetImpl()->DelFD(m_nSocketFD))
                {
                    m_nSettedEvents = 0;
                }
            }
            CloseSocket(m_nSocketFD);
        }
        m_nSocketFD = INVALID_SOCKET;
        if (IsConnecting())
        {
            if (m_cbConnectTimeoutTimer)
            {
                m_cbConnectTimeoutTimer();
                m_cbConnectTimeoutTimer = nullptr;
            }
            STaskProcessor<ENetErrorCode> sTaskProcessor = m_sConnectTask;
            m_sConnectTask = STaskProcessor<ENetErrorCode>();
            sTaskProcessor.Resolve(ENetErrorCode::OperationAborted);
            return;
        }
        while (!m_oSendTaskQueue.empty())
        {
            SSendInfo &rsSendInfo = m_oSendTaskQueue.front();
            rsSendInfo.m_oProcessor.Resolve(ENetErrorCode::OperationAborted);
            m_oSendTaskQueue.pop_front();
        }
        while (!m_oRecvTaskQueue.empty())
        {
            SRecvInfo& rsRecvInfo = m_oRecvTaskQueue.front();
            rsRecvInfo.m_oProcessor.Resolve({ ENetErrorCode::OperationAborted, 0 });
            m_oRecvTaskQueue.pop_front();
        }
	}

    bool CTcpSocketImpl::IsConnecting() const
    {
        return m_sConnectTask && !m_sConnectTask.GetTask().IsCompleted();
    }

    bool CTcpSocketImpl::IsConnected() const
    {
        return m_sConnectTask && m_sConnectTask.GetTask().IsCompleted() && m_sConnectTask.GetTask().GetResult() == ENetErrorCode::Success;
    }

    bool CTcpSocketImpl::SetNoDelay(bool)
    {
        if (IsClosed())
            return false;
        int flag = 1;
        return 0 == setsockopt(m_nSocketFD, IPPROTO_TCP, TCP_NODELAY, (char*)&flag, sizeof(flag));
    }

    bool CTcpSocketImpl::SetSendBuffer(UINT32 dwSize)
    {
        if (IsClosed())
            return false;
        return 0 == setsockopt(m_nSocketFD, SOL_SOCKET, SO_SNDBUF, (char*)&dwSize, sizeof(dwSize));
    }

    bool CTcpSocketImpl::SetRecvBuffer(UINT32 dwSize)
    {
        if (IsClosed())
            return false;
        return 0 == setsockopt(m_nSocketFD, SOL_SOCKET, SO_RCVBUF, (char*)&dwSize, sizeof(dwSize));
    }

    bool CTcpSocketImpl::_AddPoll(bool bOut)
    {
        if (IsClosed())
            return false;
        auto poNetService = m_poNetService.lock();
        if (!poNetService)
            return false;
        auto wBackup = m_nSettedEvents;
        if (m_nSettedEvents != 0)
        {
            if (bOut)
            {
#if _WIN32
                if (((m_nSettedEvents & POLLOUT) != POLLOUT))
                {
                    m_nSettedEvents |= POLLOUT;
#else
                if (((m_nSettedEvents & EPOLLOUT) != EPOLLOUT))
                {
                    m_nSettedEvents |= EPOLLOUT | EPOLLERR | EPOLLET;
#endif
                    if (!poNetService->GetImpl()->ModFD(m_nSocketFD, m_nSettedEvents))
                    {
                        m_nSettedEvents = wBackup;
                        return false;
                    }
                }
                return true;
            }
            else
            {
#if _WIN32
                if (((m_nSettedEvents & POLLIN) != POLLIN))
                {
                    m_nSettedEvents |= POLLIN;
#else
                if (((m_nSettedEvents & EPOLLIN) != EPOLLIN))
                {
                    m_nSettedEvents |= EPOLLIN | EPOLLERR | EPOLLET;
#endif
                    if (!poNetService->GetImpl()->ModFD(m_nSocketFD, m_nSettedEvents))
                    {
                        m_nSettedEvents = wBackup;
                        return false;
                    }
                }
                return true;
            }
        }
        else
        {
#if _WIN32
            m_nSettedEvents |= (bOut ? POLLOUT : POLLIN);
#else
            m_nSettedEvents |= (bOut ? EPOLLOUT : EPOLLIN) | EPOLLERR | EPOLLET;
#endif
            auto poWeakPtr = m_poSelf;
#if _WIN32
            bool bAddFDRet = poNetService->GetImpl()->AddFD(m_nSocketFD, m_nSettedEvents, [poWeakPtr](SOCKET, SHORT wReturnEvents)
#else
            bool bAddFDRet = m_poNetService->GetImpl()->AddFD(m_nSocketFD, m_nSettedEvents, [poWeakPtr](SOCKET, UINT32 wReturnEvents)
#endif
            {
                auto poSelf = poWeakPtr.lock();
                if (poSelf == nullptr || poSelf->IsClosed())
                    return;
                CTcpSocketImpl& roSelf = *poSelf;
#if _WIN32
                if ((wReturnEvents & POLLERR) != 0)
#else
                if ((wReturnEvents & EPOLLERR) != 0)
#endif
                {
                    ENetErrorCode eNetErrorCode = GetLastError();
                    if (roSelf.IsConnecting())
                    {
                        auto sTaskProcessor = roSelf.m_sConnectTask;
                        roSelf.m_sConnectTask = STaskProcessor<ENetErrorCode>();
                        sTaskProcessor.Resolve(eNetErrorCode);
                    }
                    else
                    {
                        while (!roSelf.m_oSendTaskQueue.empty())
                        {
                            SSendInfo& rsSendInfo = roSelf.m_oSendTaskQueue.front();
                            rsSendInfo.m_oProcessor.Resolve(eNetErrorCode);
                            roSelf.m_oSendTaskQueue.pop_front();
                        }
                        while (!roSelf.m_oRecvTaskQueue.empty())
                        {
                            SRecvInfo& rsRecvInfo = roSelf.m_oRecvTaskQueue.front();
                            rsRecvInfo.m_oProcessor.Resolve({ eNetErrorCode, 0 });
                            roSelf.m_oRecvTaskQueue.pop_front();
                        }
                    }
                    return;
                }
#if _WIN32
                if ((wReturnEvents & POLLOUT) != 0)
#else
                if ((wReturnEvents & EPOLLOUT) != 0)
#endif
                {
                    if (roSelf.IsConnecting())
                    {
                        roSelf._RemovePollOut();
                        if (roSelf.m_cbConnectTimeoutTimer)
                        {
                            roSelf.m_cbConnectTimeoutTimer();
                            roSelf.m_cbConnectTimeoutTimer = nullptr;
                        }
                        roSelf.m_sConnectTask.Resolve(ENetErrorCode::Success);
                    }
                    else
                    {
                        roSelf._Send();
                    }
                }
#if _WIN32
                if ((wReturnEvents & POLLIN) != 0)
#else
                if ((wReturnEvents & EPOLLIN) != 0)
#endif
                {
                    roSelf._Recv();
                }
            });
            if (!bAddFDRet)
            {
                m_nSettedEvents = wBackup;
                return false;
            }
            return true;
        }
    }

    bool CTcpSocketImpl::_RemovePollOut()
    {
        if (IsClosed())
            return false;
        auto poNetService = m_poNetService.lock();
        if (!poNetService)
            return false;
#if _WIN32
        if ((m_nSettedEvents & POLLOUT) == 0)
            return true;
        auto wBackup = m_nSettedEvents;
        m_nSettedEvents &= ~POLLOUT;
#else
        if ((m_nSettedEvents & EPOLLOUT) == 0)
            return true;
        auto wBackup = m_nSettedEvents;
        m_nSettedEvents &= ~EPOLLOUT;
#endif
            
        if (m_nSettedEvents == 0)
        {
            if (!poNetService->GetImpl()->DelFD(m_nSocketFD))
            {
                m_nSettedEvents = wBackup;
                return false;
            }
            return true;
        }
        if (!poNetService->GetImpl()->ModFD(m_nSocketFD, m_nSettedEvents))
        {
            m_nSettedEvents = wBackup;
            return false;
        }
        return true;
    }

    bool CTcpSocketImpl::_RemovePollIn()
    {
        if (IsClosed())
            return false;
        auto poNetService = m_poNetService.lock();
        if (!poNetService)
            return false;
#if _WIN32
        if ((m_nSettedEvents & POLLIN) == 0)
            return true;
        auto wBackup = m_nSettedEvents;
        m_nSettedEvents &= ~POLLIN;
#else
        if ((m_nSettedEvents & EPOLLIN) == 0)
            return true;
        auto wBackup = m_nSettedEvents;
        m_nSettedEvents &= ~EPOLLIN;
#endif
        if (m_nSettedEvents == 0)
        {
            if (!poNetService->GetImpl()->DelFD(m_nSocketFD))
            {
                m_nSettedEvents = wBackup;
                return false;
            }
            return true;
        }
        if (!poNetService->GetImpl()->ModFD(m_nSocketFD, m_nSettedEvents))
        {
            m_nSettedEvents = wBackup;
            return false;
        }
        return true;
    }
}